from libsvm.commonutil import svm_read_problem
from libsvm.svmutil import svm_train, svm_predict, svm_problem, svm_parameter


class TFPN():
    def __init__(self, labels, p_labels):
        self.tp = self.fn = self.fp = self.tn = 0

        for pair in zip(labels, p_labels):
            if pair == (1, 1):
                self.tp += 1
            elif pair == (1, -1):
                self.fn += 1
            elif pair == (-1, -1):
                self.tn += 1
            else:
                self.fp += 1
        self.t = self.tp + self.tn
        self.f = self.fp + self.fn
        self.all = self.t + self.f
        self.p_all = self.tp + self.fp
        self.r_all = self.tp + self.fn
        self.ACC = self.t / self.all
        self.P = self.tp / self.p_all
        self.R = self.tp / self.r_all
        self.F1 = 2 * self.P * self.R / (self.P + self.R)

    def __str__(self):
        return str((self.ACC, self.P, self.R, self.F1))


def check_problem(file_name):
    test_file = '{}.t'.format(file_name)
    y, x = svm_read_problem(file_name)
    y_test, x_test = svm_read_problem(test_file)
    problem = svm_problem(y, x)
    parameter = svm_parameter('-t 0 -c 1 -q')
    model = svm_train(problem, parameter)
    labels_test, *_ = svm_predict(y_test, x_test, model)
    labels, *_ = svm_predict(y, x, model)
    print(file_name, TFPN(y_test, labels_test))
    print(test_file, TFPN(y, labels))


if __name__ == '__main__':
    check_problem('a1a')
    check_problem('a2a')
